import jax
import jax.numpy as jnp
from jax.scipy.signal import convolve2d
from jax.scipy.ndimage import map_coordinates
import time

jax.config.update('jax_platform_name', 'cpu')  # Ensure JAX uses GPU by default

class BarotropicModel:
    """
    A simple barotropic vorticity equation model for atmospheric flow using JAX, optimized for GPU.
    This model represents non-divergent 2D flow on a beta-plane.
    """
    def __init__(self, nx=20, ny=20, dx=100000, dy=100000, dt=1200, 
                 beta=2e-11, diffusion=2000):
        self.nx = nx
        self.ny = ny
        self.state_size = nx * ny
        self.dx = dx
        self.dy = dy
        self.dt = dt
        self.beta = beta
        self.diffusion = diffusion
        
        # Initialize timings dictionary
        self.timings = {}
        
        self._setup_grid()
        self._setup_diffusion_kernel()  # Precompute diffusion kernel
        
    def _setup_grid(self):
        """Setup spatial grids on GPU"""
        self.x = jax.device_put(jnp.linspace(0, self.nx * self.dx, self.nx))
        self.y = jax.device_put(jnp.linspace(0, self.ny * self.dy, self.ny))
        self.X, self.Y = jax.device_put(jnp.meshgrid(self.x, self.y, indexing='xy'))
        
        # FFT grids with proper indexing
        kx = jnp.fft.fftfreq(self.nx, d=self.dx/(2*jnp.pi))
        ky = jnp.fft.fftfreq(self.ny, d=self.dy/(2*jnp.pi))
        self.KX, self.KY = jax.device_put(jnp.meshgrid(kx, ky, indexing='ij'))
        self.K2 = jax.device_put(self.KX**2 + self.KY**2).at[0,0].set(1.0)
        
    def _setup_diffusion_kernel(self):
        """Precompute Gaussian diffusion kernel on GPU"""
        sigma = jnp.sqrt(2*self.diffusion*self.dt)/self.dx
        kernel_size = int(2*jnp.ceil(2*sigma)+1)
        ax = jax.device_put(jnp.linspace(-(kernel_size-1)/2, (kernel_size-1)/2, kernel_size))
        xx, yy = jnp.meshgrid(ax, ax)
        kernel = jnp.exp(-0.5*(xx**2 + yy**2)/sigma**2)
        self.diffusion_kernel = jax.device_put(kernel/kernel.sum())

    def _reshape_state(self, x_flat):
        return x_flat.reshape(self.ny, self.nx)
    
    def _flatten_state(self, x_2d):
        return x_2d.flatten()
    
    def _compute_streamfunction(self, vorticity):
        """Spectral Poisson solver"""
        vort_fft = jnp.fft.fft2(vorticity)
        psi_fft = -vort_fft / self.K2
        return jnp.real(jnp.fft.ifft2(psi_fft.at[0,0].set(0)))
    
    def _compute_velocity(self, psi):
        u = -jnp.gradient(psi, self.dy, axis=0)
        v = jnp.gradient(psi, self.dx, axis=1)
        return u, v
    
    def _apply_diffusion(self, vorticity):
        """GPU-optimized diffusion application"""
        pad = self.diffusion_kernel.shape[0]//2
        padded = jnp.pad(vorticity, pad, mode='wrap')
        convolved = convolve2d(padded, self.diffusion_kernel, mode='same')
        return convolved[pad:-pad, pad:-pad]
    
    def step(self, x):
        """JIT-compiled time step with GPU-optimized advection"""
        start_time = time.time()
        
        vorticity = self._reshape_state(x)
        
        # Calculate streamfunction from vorticity
        t0 = time.time()
        psi = self._compute_streamfunction(vorticity)
        psi = jax.block_until_ready(psi)
        self.timings['compute_streamfunction'] = self.timings.get('compute_streamfunction', 0) + (time.time() - t0)
        
        # Calculate velocity field
        t0 = time.time()
        u, v = self._compute_velocity(psi)
        u, v = jax.block_until_ready(u), jax.block_until_ready(v)
        self.timings['compute_velocity'] = self.timings.get('compute_velocity', 0) + (time.time() - t0)
        
        # Semi-Lagrangian advection with GPU interpolation
        t0 = time.time()
        Xd = (self.X - u*self.dt) / self.dx
        Yd = (self.Y - v*self.dt) / self.dy
        vort_new = map_coordinates(vorticity, jnp.stack([Yd, Xd]), 1, 'wrap')
        vort_new = jax.block_until_ready(vort_new)
        self.timings['interpolation'] = self.timings.get('interpolation', 0) + (time.time() - t0)
        
        # Physical processes
        vort_new += self.beta * v * self.dt
        
        # Apply diffusion
        t0 = time.time()
        vort_new = self._apply_diffusion(vort_new)
        vort_new = jax.block_until_ready(vort_new)
        self.timings['diffusion'] = self.timings.get('diffusion', 0) + (time.time() - t0)
        
        result = self._flatten_state(vort_new)
        result = jax.block_until_ready(result)
        
        self.timings['step_total'] = self.timings.get('step_total', 0) + (time.time() - start_time)
        return result
    
    def integrate(self, x0, t_steps):
        """GPU-optimized integration"""
        start_time = time.time()
        
        # Create an un-JIT compiled version for timing individual operations
        if 'step' not in self.timings:
            step_result = self.step(x0)
            jax.block_until_ready(step_result)
        
        # Now create the JIT-compiled version for actual integration
        @jax.jit
        def step_function(x):
            return self.step(x)
            
        @jax.jit
        def scan_step(carry, _):
            next_state = step_function(carry)
            return next_state, next_state
            
        _, trajectory = jax.lax.scan(scan_step, x0, jnp.arange(t_steps))
        result = jnp.vstack([x0[None, ...], trajectory])
        
        # Ensure computation is complete before timing
        jax.block_until_ready(result)
        end_time = time.time()
        self.timings['integrate'] = self.timings.get('integrate', 0) + (end_time - start_time)
        
        return result
    
    # Remaining methods use same GPU optimizations
    def tangent_linear(self, x_ref):
        x_ref = jax.device_put(x_ref)
        def body_fn(i, M):
            eps = 1e-6
            dx = jax.device_put(jnp.zeros_like(x_ref).at[i].set(eps))
            return M.at[:,i].set((self.step(x_ref+dx)-self.step(x_ref))/eps)
        return jax.lax.fori_loop(0, len(x_ref), body_fn, jnp.zeros((len(x_ref), len(x_ref))))
    
    def adjoint(self, x_ref):
        return self.tangent_linear(x_ref).T
    
    def generate_random_state(self, mean_amplitude=1.0, seed=None):
        key = jax.random.PRNGKey(seed if seed is not None else 0)
        # GPU-optimized spectral initialization
        k = jnp.sqrt(self.KX**2 + self.KY**2)
        spectrum = (jnp.exp(2j*jnp.pi*jax.random.uniform(key, self.KX.shape)) *
                    (k > 1) * (k <= min(self.nx,self.ny)//2) * k**(-1.5))
        return self._flatten_state(jnp.real(jnp.fft.ifft2(spectrum)))
    
    def get_timings(self):
        """Return the timing statistics for the model functions"""
        return self.timings
        
    def print_timing_statistics(self):
        """Print detailed statistics about function execution times"""
        if not self.timings:
            print("No timing data available.")
            return
            
        # Calculate total model execution time excluding the integrate time
        # (since integrate includes the sum of all step operations)
        operation_times = {k: v for k, v in self.timings.items() if k != 'integrate'}
        total_time = sum(operation_times.values())
        
        print("\n===== Timing Statistics =====")
        print(f"Total execution time: {total_time:.6f} seconds")
        
        # Sort operations by time (descending)
        sorted_ops = sorted(operation_times.items(), key=lambda x: x[1], reverse=True)
        
        print("\nBreakdown by operation:")
        print("-----------------------")
        for op, time_spent in sorted_ops:
            percentage = (time_spent / total_time) * 100
            print(f"{op:<25}: {time_spent:.6f} s ({percentage:.2f}%)")
            
        # If we have step_total in the timings
        if 'step_total' in self.timings:
            # Estimate number of steps performed
            step_count = int(self.timings['step_total'] / self.timings.get('step', self.timings['step_total']))
            avg_step_time = self.timings['step_total'] / step_count if step_count > 0 else 0
            print(f"\nAverage time per time step: {avg_step_time:.6f} seconds")